import sys
sys.path.append('..')
# sys.path.append(r'/home/user10000630/notespace/Chatbot_227_ef9a')
from retrieve.hnsw import HnswSelector, transform_by_embedding
from site_packages.utils.job import DataOp
import jieba
from site_packages.ml_libs.nlp.stopwords import Stopwords
from business.intention import Intention
import logging
from ranking.ranker import Ranker
from generative.predict import bertSeq2Seq
from configs.settings import SAVED_GENERATIVE_PATH, DEVICE
import os
logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class BakaRobot:
    """
    接收用户请求，完成意图识别、检索、返回答案、闲聊的流程
    args:
        rank_model_type: logistic或者lgbm
    attr:
        hnsw: 训练好的hnsw检索器
        hnsw_corpus:用于结合hnsw返回的结果索引，返回给用户文字答案
        stopwords: 停用词过滤
        ranker: 精排
        talking_bert: 用于闲聊的bert
    """
    docvec_dim = 100
    
    def __init__(self, rank_model_type='logistic'):
        self.intention = Intention(training=False)
        self.hnsw = HnswSelector(dim=self.docvec_dim, layer_num=100, gpu=True, neighbor_num=64, training=False)
        self.hnsw_corpus = DataOp.load_data('dev')['customer']
        self.stopwords = Stopwords()
        self.ranker = Ranker(training=False, model_type=rank_model_type)
        self.talking_bert = bertSeq2Seq(os.path.join(SAVED_GENERATIVE_PATH, 'bert.model.epoch.29'), DEVICE)

    def format_query(self, query):
        query_list = jieba.lcut(query)
        query_list = self.stopwords.clean(query_list)
        return query_list
    
    def find_nearest_queries(self, query_list):
        docvec = transform_by_embedding([query_list])
        distances, indices = self.hnsw.search(docvec)
        results = self.hnsw_corpus[indices[0]]
        return [' '.join(result) for result in results]
            
    def answer(self, query):
        # 判断是否为业务方面的提问
        fit_this, prob = self.intention(text=query)
        if fit_this:
            logger.info("这个问题是业务类提问")
            # 搜索相似问题：
            query_list = self.format_query(query)
            similar_queries = self.find_nearest_queries(query_list)
            ranking_result = self.ranker.predict(' '.join(query_list), similar_queries)
            return ranking_result
        else:
            logger.info("这个问题是闲聊类提问")
            return self.talking_bert.generate(query, k=4)
    
    
if __name__ == "__main__":
    query = "你好"   
    robot = BakaRobot()
    ranking_result = robot.answer(query)
    print("ranking_result:", ranking_result)